"""
页面模块 - 包含首页和详情页的显示功能
"""

import streamlit as st
import os
import base64
from datetime import datetime

from modules.ai_client import generate_response
from modules.embeddings import embed_text, search_similar_docs, get_document_for_artifact
from modules.knowledge_graph import get_related_kg_info, format_kg_info, visualize_artifact_kg, create_artifact_info_card
from modules.timeline import create_timeline_chart
from modules.image_utils import save_uploaded_image, classify_artifact_with_ernie, get_artifact_sample_images
from modules.prompts import generate_quick_queries

def show_home_page(client, embedding_model, kg_data, index, documents, artifact_images):
    """显示首页"""
    st.markdown('<h1 class="main-title">🐯 智慧动物园讲解助手</h1>', unsafe_allow_html=True)
    
    if client is None:
        st.warning("⚠️ 注意: ERNIE客户端未初始化，文物智能问答功能不可用。其他功能可以正常使用。请在侧边栏设置API密钥。")
    
    tab1, tab2, tab3 = st.tabs(["动物图片识别", "动物知识问答", "尚未开发"])
    
    with tab1:
        st.markdown("""
        <div class="upload-title-container">
            <div class="upload-title">上传动物图片，AI将自动识别并讲解</div>
            <div class="upload-subtitle">支持各类动物，如狮子、老虎、大象等</div>
        </div>
        """, unsafe_allow_html=True)
        
        upload_col, result_col = st.columns(2)
        
        with upload_col:
            uploaded_file = st.file_uploader(
                "上传动物图片", 
                type=["jpg", "jpeg", "png"],
                label_visibility="collapsed",
                accept_multiple_files=False
            )
            
            if uploaded_file:
                file_size = len(uploaded_file.getvalue())
                max_size = 10 * 1024 * 1024  # 10MB
                
                if file_size > max_size:
                    st.error(f"文件过大！请上传小于10MB的图片。当前文件大小: {file_size/1024/1024:.2f}MB")
                else:
                    image_data = uploaded_file.getvalue()
                    image_b64 = base64.b64encode(image_data).decode()
                    
                    st.markdown(f"""
                    <div class="uploaded-image-container">
                        <img src="data:image/jpeg;base64,{image_b64}" class="uploaded-image" alt="上传的文物图片">
                        <div class="image-caption">上传的动物图片</div>
                    </div>
                    """, unsafe_allow_html=True)
        
        with result_col:
            if uploaded_file and len(uploaded_file.getvalue()) <= 10 * 1024 * 1024:
                image_path = save_uploaded_image(uploaded_file)
                
                if image_path:
                    with st.spinner("🔍 正在识别动物..."):
                        artifact_name, confidence = classify_artifact_with_ernie(client, image_path, kg_data)
                        
                        if confidence > 0.5:
                            st.markdown(f"""
                            <div class="recognition-result">
                                <div class="recognition-title">📷 识别结果</div>
                                <p><b>动物名称：</b>{artifact_name}</p>
                                <p><b>识别置信度：</b>{confidence:.2f}</p>
                            </div>
                            """, unsafe_allow_html=True)
                            
                            st.session_state.current_artifact = artifact_name
                            
                            artifact_doc = get_document_for_artifact(artifact_name, documents)
                            
                            relations = get_related_kg_info(artifact_name, kg_data)
                            kg_info = format_kg_info(relations)
                            
                            context = f"{artifact_doc}\n\n{kg_info}" if artifact_doc else kg_info
                            
                            intro_query = f"请用50字左右简要介绍{artifact_name}的特点和习性。"
                            
                            with st.spinner("生成动物介绍中..."):
                                if client:
                                    brief_intro = generate_response(client, intro_query, context, None, artifact_name, stream_output=False)

                                    if len(brief_intro) > 300:
                                        brief_intro = brief_intro[:297] + "..."
                                else:
                                    brief_intro = f"无法生成{artifact_name}的介绍，ERNIE客户端未初始化。"
                            
                            st.markdown(f"""
                            <div class="recognition-result" style="background-color: rgba(48, 140, 108, 0.1); border-left-color: #308C6C;">
                                <div class="recognition-title">📚 动物简介</div>
                                <p>{brief_intro}</p>
                            </div>
                            """, unsafe_allow_html=True)
                            
                            if st.button(f"查看《{artifact_name}》详细讲解", key="view_details"):
                                st.session_state.page = "detail"
                                st.rerun()
                        else:
                            st.warning("无法确定图片中的动物，请尝试不同的图片或直接从侧边栏选择动物。")
                else:
                    st.error("图片保存失败，请重试。")
            else:
                st.markdown("""
                <div class="result-area">
                    <img src="https://img.icons8.com/fluency/96/camera.png" style="width: 48px; height: 48px; margin-bottom: 15px;">
                    <p style="color: #666; font-size: 1.1rem;">请在左侧上传动物图片</p>
                    <p style="color: #888; font-size: 0.9rem;">AI将自动识别并为您提供专业讲解</p>
                </div>
                """, unsafe_allow_html=True)
        
        if not uploaded_file:
            st.markdown("<h3 class='sub-title' style='text-align: center; margin-top: 30px;'>动物探索</h3>", unsafe_allow_html=True)
            
            featured_artifacts = [
                "狮子", "老虎", "大象", "长颈鹿", "斑马", "企鹅"
            ]
            
            row1_cols = st.columns(3)
            
            for i, col in enumerate(row1_cols):
                if i < len(featured_artifacts):
                    artifact_name = featured_artifacts[i]
                    if artifact_name in artifact_images:
                        with col:
                            st.markdown(f"""
                            <div>
                                <img src="{artifact_images[artifact_name]}" alt="{artifact_name}" class="artifact-gallery-image">
                                <div class="artifact-gallery-caption">{artifact_name}</div>
                            </div>
                            """, unsafe_allow_html=True)
                            
                            if st.button(f"了解《{artifact_name}》", key=f"feat_{i}"):
                                st.session_state.current_artifact = artifact_name
                                st.session_state.page = "detail"
                                st.rerun()
            
            row2_cols = st.columns(3)
            
            for i, col in enumerate(row2_cols):
                idx = i + 3 
                if idx < len(featured_artifacts):
                    artifact_name = featured_artifacts[idx]
                    if artifact_name in artifact_images:
                        with col:
                            st.markdown(f"""
                            <div>
                                <img src="{artifact_images[artifact_name]}" alt="{artifact_name}" class="artifact-gallery-image">
                                <div class="artifact-gallery-caption">{artifact_name}</div>
                            </div>
                            """, unsafe_allow_html=True)
                            
                            if st.button(f"了解《{artifact_name}》", key=f"feat_{idx}"):
                                st.session_state.current_artifact = artifact_name
                                st.session_state.page = "detail"
                                st.rerun()

    
    with tab2:        
        st.markdown('<div class="highlight-box">您可以直接提问关于动物的问题，如"狮子有哪些习性？"、"企鹅的生活特点是什么？"</div>', unsafe_allow_html=True)
        
        st.markdown('<div style="display: flex; flex-wrap: wrap; gap: 10px; margin-bottom: 15px;">', unsafe_allow_html=True)
        for question in ["狮子有哪些习性？", "企鹅的生活特点是什么？", "长颈鹿的栖息地在哪", "老虎以什么为食？"]:
            st.markdown(f'<button class="quick-query-button">{question}</button>', unsafe_allow_html=True)
        st.markdown('</div>', unsafe_allow_html=True)
        
        query = st.text_input("您想了解什么？", key="query_input_home")
        
        if st.button("发送问题", key="send_query_home") and query:
            with st.spinner("🔍 正在查询相关知识..."):
                query_embedding = embed_text(query, embedding_model)
                
                search_results = search_similar_docs(query_embedding, index, documents)
                
                context = ""
                for result in search_results:
                    context += result["content"] + "\n\n"
                
                artifact_name = None
                for node in kg_data['nodes']:
                    if node['type'] == 'animal' and node['name'] in query:
                        artifact_name = node['name']
                        break
                
                kg_info = ""
                if artifact_name:
                    relations = get_related_kg_info(artifact_name, kg_data)
                    kg_info = format_kg_info(relations)
                    context += "\n" + kg_info

                response = generate_response(client, query, context, None, artifact_name, stream_output=True)
                
                st.session_state.chat_history.append({
                    "query": query,
                    "response": response,
                    "artifact": artifact_name,
                    "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                })
                
                if artifact_name:
                    col1, col2, col3 = st.columns([1, 1, 1])
                    with col2:
                        if st.button(f"查看《{artifact_name}》详情", key=f"view_details_{artifact_name}_{len(st.session_state.chat_history)}"):
                            st.session_state.current_artifact = artifact_name
                            st.session_state.page = "detail"
                            st.rerun()
        
        if st.session_state.chat_history:
            st.markdown('<h3 class="sub-title">最近问答</h3>', unsafe_allow_html=True)
            
            for i, exchange in enumerate(st.session_state.chat_history[-3:]):
                with st.expander(f"问题：{exchange['query']}", expanded=False):
                    st.markdown(f"<div class='history-response'>{exchange['response']}</div>", unsafe_allow_html=True)
                    st.markdown(f"<div style='text-align:right;font-size:0.8rem;color:#666666;'>{exchange['timestamp']}</div>", unsafe_allow_html=True)
        
        st.markdown('</div>', unsafe_allow_html=True)
    
    with tab3:        
        timeline_fig = create_timeline_chart()
        st.plotly_chart(timeline_fig)
        
        st.markdown('<h3 class="sub-title">历史朝代文物概览</h3>', unsafe_allow_html=True)
        
        dynasty_artifacts = {}
        from modules.timeline import get_artifact_dynasty_mapping
        artifact_dynasty_mapping = get_artifact_dynasty_mapping()
        
        for artifact, dynasties in artifact_dynasty_mapping.items():
            for dynasty in dynasties:
                if dynasty not in dynasty_artifacts:
                    dynasty_artifacts[dynasty] = []
                dynasty_artifacts[dynasty].append(artifact)
        
        from modules.timeline import get_chinese_dynasties
        for dynasty, artifacts in sorted(dynasty_artifacts.items(), key=lambda x: get_chinese_dynasties()[0]['start'] if x[0] not in [d['name'] for d in get_chinese_dynasties()] else next((d['start'] for d in get_chinese_dynasties() if d['name'] == x[0]), 0)):
            with st.expander(f"{dynasty} 时期文物"):
                artifact_links = []
                for artifact in artifacts:
                    artifact_links.append(f'<a href="#" onclick="parent.postMessage({{action: \'select_artifact\', name: \'{artifact}\'}}, \'*\')">{artifact}</a>')
                
                st.markdown(f"<div>{'、'.join(artifact_links)}</div>", unsafe_allow_html=True)
                
                col1, col2, col3 = st.columns([1, 3, 1])
                with col2:
                    if st.button(f"探索{dynasty}时期文物", key=f"dynasty_{dynasty}"):
                        if artifacts:
                            st.session_state.current_artifact = artifacts[0]
                            st.session_state.page = "detail"
                            st.rerun()

def show_detail_page(client, embedding_model, kg_data, index, documents, artifact_images):
    """显示详情页"""
    artifact_name = st.session_state.current_artifact
    if not artifact_name:
        st.error("未找到动物信息")
        if st.button("返回首页"):
            st.session_state.page = "home"
            st.rerun()
        return
    
    col_title, col_back = st.columns([5, 1])
    
    with col_title:
        st.markdown(f'<h2 class="sub-title">《{artifact_name}》详情</h2>', unsafe_allow_html=True)
    
    with col_back:
        if st.button("返回 →", key="back_button"):
            st.session_state.page = "home"
            st.rerun()
    
    image_col, info_col = st.columns([3, 2])
    
    with image_col:
        if artifact_name in artifact_images:
            st.markdown(f"""
            <div style="background-color: #f5f5f5; border-radius: 8px; padding: 10px; text-align: center;">
                <img src="{artifact_images[artifact_name]}" alt="{artifact_name}" 
                     style="max-width: 100%; max-height: 400px; object-fit: contain; border-radius: 4px;">
            </div>
            """, unsafe_allow_html=True)
        else:
            st.info("暂无动物图片")
            st.markdown("""
            <div style="background-color: #f0f0f0; height: 400px; display: flex; justify-content: center; align-items: center; border-radius: 8px;">
                <p style="color: #666;">暂无图片</p>
            </div>
            """, unsafe_allow_html=True)
    
    with info_col:
        st.markdown('## 📌 动物基本信息', unsafe_allow_html=True)
        
        info_card = create_artifact_info_card(artifact_name, kg_data)
        
        if info_card:
            for key, value in info_card.items():
                st.markdown(f"**{key}:** {value}")
        else:
            st.markdown("暂无详细信息")
    
    st.markdown('<div style="height: 30px;"></div>', unsafe_allow_html=True)  # 间隔
    
    tabs = st.tabs(["文物详细介绍", "知识图谱", "暂未开发", "快速提问"])
    
    with tabs[0]:
        artifact_doc = get_document_for_artifact(artifact_name, documents)
        
        if 'artifact_intro' not in st.session_state or st.session_state.get('current_artifact_intro') != artifact_name:
            with st.spinner():
                spinner_container = st.empty()
                spinner_container.markdown("""
                <div class="simple-spinner">
                    <div class="spinner-circle"></div>
                </div>
                """, unsafe_allow_html=True)
                
                intro_query = f"请全面介绍{artifact_name}的习性、栖息地、食物构成和重要意义。"
                
                relations = get_related_kg_info(artifact_name, kg_data)
                kg_info = format_kg_info(relations)
                
                context = f"{artifact_doc}\n\n{kg_info}" if artifact_doc else kg_info
                
                if client:
                    intro_response = generate_response(client, intro_query, context, None, artifact_name, stream_output=False)
                    spinner_container.empty()
                else:
                    intro_response = f"**{artifact_name}**\n\n{artifact_doc}\n\n{kg_info}\n\n(注：ERNIE API未设置，无法生成更详细的讲解)"
                    spinner_container.empty()
                
                st.session_state.artifact_intro = intro_response
                st.session_state.current_artifact_intro = artifact_name
        
        st.markdown(f'<div class="response-container">{st.session_state.artifact_intro}</div>', unsafe_allow_html=True)
    
    with tabs[1]:
        kg_fig = visualize_artifact_kg(artifact_name, kg_data)
        if kg_fig:
            st.plotly_chart(kg_fig)
        else:
            st.info("无法生成知识图谱可视化")
    
    #with tabs[2]:
    #    timeline_fig = create_timeline_chart(artifact_name)
    #    st.plotly_chart(timeline_fig)
    
    with tabs[3]:
        quick_queries = generate_quick_queries(artifact_name)
        
        st.markdown('<h3 class="sub-title">快速提问</h3>', unsafe_allow_html=True)
        
        cols = st.columns(3)
        
        for i, query in enumerate(quick_queries):
            col_index = i % 3
            with cols[col_index]:
                if st.button(query, key=f"quick_q_{i}", help="点击提问"):
                    st.session_state.current_query = query
        
        st.markdown('<h3 class="sub-title">自定义提问</h3>', unsafe_allow_html=True)
        custom_query = st.text_input(f"关于《{artifact_name}》，您想了解什么？", key="custom_query_detail")
        
        if hasattr(st.session_state, 'current_query'):
            custom_query = st.session_state.current_query
            del st.session_state.current_query
        
        if st.button("提问", key="send_custom_query") or custom_query:
            if custom_query:
                with st.spinner("生成回答中..."):
                    relations = get_related_kg_info(artifact_name, kg_data)
                    kg_info = format_kg_info(relations)
                    
                    context = f"{artifact_doc}\n\n{kg_info}" if artifact_doc else kg_info
                    
                    response = generate_response(client, custom_query, context, None, artifact_name, stream_output=True)
                    
                    st.session_state.chat_history.append({
                        "query": custom_query,
                        "response": response,
                        "artifact": artifact_name,
                        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                    })
                    
            else:
                st.warning("请输入您的问题") 